from utils import print_timing
from pheno import Pheno_Base
from anno import Anno_Affy6, Anno_Illumina
from geno import Geno_Single_Affy6, Geno_Single_Illumina, Geno_Shard_Affy6, Geno_Shard_Illumina

class Platform(object):
    def __init__(self, pheno, anno, geno, schema, dbtype, dbstring, tmpdir, table_dict, jobs):
        self.pheno = pheno
        self.anno = anno
        self.geno = geno
        self.schema = schema
        self.dbtype = dbtype
        self.dbstring = dbstring
        self.tmpdir = tmpdir
        self.table_dict = table_dict
        self.jobs = jobs

    @print_timing
    def load_dataset(self):
        import os, sys
        from utils import file_not_exist, safe_mkdir
        from dbutils import get_db_type
        from dbfunctions import create_drop_constraint_if_exists_function, create_or_replace_plpgsql, create_decode_genotype_function
        import shutil
        from utils import create_engine_wrapper
        from sqlalchemy.orm import mapper, relation, sessionmaker
        import logging
        logging.basicConfig()
        dbtype = get_db_type(self.dbstring)
        if dbtype != self.dbtype:
            sys.exit("error: this database has type %s, but must have type '%s'"%(dbtype, self.dbtype))
        # Exit if any of patientdatafilenames, annotfilename, callsfilename don't exist
        if [x for x in self.pheno.phenofiles + [self.anno.annotfile, self.geno.genofile] if file_not_exist(x)]:
            sys.exit("Exiting in Platform.load_dataset in plat.py...")
        print "making %s tables corresponding to schema %s and dbstring %s"%(self.dbtype, self.schema, self.dbstring)
        safe_mkdir(self.tmpdir)
        td = self.table_dict
        Allele, Chromo, Pheno, Idlink, Snpval = td['Allele'], td['Chromo'], td['Pheno'], td['Idlink'], td['Snpval'],
        db = create_engine_wrapper(self.dbstring)
        meta = self.table_dict['metadata']
        meta.bind = db
        meta.create_all()
        Session = sessionmaker()
        session = Session(bind=db)
        session.execute(create_or_replace_plpgsql)
        session.execute(create_drop_constraint_if_exists_function)
        session.execute(create_decode_genotype_function)
        sexd = self.pheno.get_sex_dict(session)
        raced = self.pheno.get_race_dict(session)

        patient_expids = []
        for phenofile in self.pheno.phenofiles:
            patient_expids += self.pheno.get_patient_expids(phenofile)

        expids = self.geno.celfilenames()

        # Pheno objects.
        if session.query(Pheno).count() == 0:
            for phenofile in self.pheno.phenofiles:
                session.add_all(self.pheno.make_pheno_table(phenofile, sexd, raced))
            session.add_all(self.pheno.dummy_patients(expids, patient_expids, sexd, raced))

        # Allele objects
        alleled = {}
        for a in ['A', 'C', 'G', 'T']:
            alleled[a] = Allele(a)
        if session.query(Allele).count() == 0:
            session.add_all(alleled.values())
        session.commit()
        session.flush()

        #Idlink objects
        if session.query(Idlink).count() == 0:
            expids =  self.geno.celfilenames()
            self.pheno.make_idlink_table(db, expids, patient_expids)
        else:
            sys.exit("idlink table already exists, exiting")
        session.commit()
        session.flush()

        # Chromo objects.
        chromolst = []
        if session.query(Chromo).count() == 0:
            for c in range(1,23):
                chromolst.append(Chromo(c, str(c)))
            chromolst.append(Chromo(23, 'X'))
            chromolst.append(Chromo(24, 'Y'))
            chromolst.append(Chromo(25, 'XY'))
            chromolst.append(Chromo(26, 'MT'))
            session.add_all(chromolst)
        session.commit()
        session.flush()
        # Anno objects.
        Anno = self.table_dict['Anno']
        if session.query(Anno).count() == 0:
            self.anno.make_anno_table(session, alleled)
        else:
            sys.exit("anno table already exists, exiting")
        # Snpval objects.
        if session.query(Snpval).count() == 0:
            svlst = []
            for i in [-1, 0, 1, 2]:
                svlst.append(Snpval(i))
            session.add_all(svlst)
        session.commit()
        session.flush()
        self.load_geno(session)
        db.dispose()
        if os.path.exists(self.tmpdir):
            shutil.rmtree(self.tmpdir)

class Single(Platform):

    def load_geno(self, session):
        Geno_table = self.table_dict['Geno_table']
        geno_count = session.query(Geno_table).count()
        session.close()
        if geno_count == 0:
            self.geno.make_geno_from_file(session)
        else:
            sys.exit("geno table already exists, exiting")

class Single_Affy6(Single):
    def __init__(self, phenofiles, annotfile, genofile, schema, dbtype, dbstring, tmpdir, jobs):
        from dbschema import make_tables_affy6
        table_dict = make_tables_affy6(schema)
        pheno = Pheno_Base(phenofiles, dbstring, table_dict)
        anno = Anno_Affy6(annotfile, schema, tmpdir, table_dict)
        geno = Geno_Single_Affy6(genofile, schema, dbstring, tmpdir, table_dict, jobs)
        Single.__init__(self, pheno, anno, geno, schema, dbtype, dbstring, tmpdir, table_dict, jobs)

class Single_Illumina(Single):
    def __init__(self, phenofiles, annotfile, genofile, schema, alleletype, dbtype, dbstring, tmpdir, jobs):
        from dbschema import make_tables_illumina
        table_dict = make_tables_illumina(schema)
        pheno = Pheno_Base(phenofiles, dbstring, table_dict)
        anno = Anno_Illumina(annotfile, schema, tmpdir, table_dict)
        geno = Geno_Single_Illumina(genofile, schema, dbstring, tmpdir, table_dict, alleletype, jobs)
        Single.__init__(self, pheno, anno, geno, schema, dbtype, dbstring, tmpdir, table_dict, jobs)

    def load_dataset(self):
        import sys
        alleletype = self.geno.alleletype
        if alleletype != "forward" and alleletype != "top":
            sys.exit("alleletype must be either 'forward' or 'top', not %s"%alleletype)
        super(Single_Illumina, self).load_dataset()

class Shard(Platform):

    def load_geno(self, session):
        # Geno_table = table_dict['Geno_table']
        # geno_count = session.query(Geno_table).count()
        # session.close()
        # if geno_count == 0:
        #     self.geno.make_geno_table_from_file(db)
        # else:
        #     sys.exit("geno table already exists, exiting")

        self.geno.make_geno_from_file(session)

class Shard_Affy6(Shard):
    def __init__(self, phenofiles, annotfile, genofile, schema, dbtype, dbstring, tmpdir, jobs):
        from dbschema import make_tables_shard_affy6
        table_dict = make_tables_shard_affy6(schema)
        pheno = Pheno_Base(phenofiles, dbstring, table_dict)
        anno = Anno_Affy6(annotfile, schema, tmpdir, table_dict)
        geno = Geno_Shard_Affy6(genofile, schema, dbstring, tmpdir, table_dict, jobs)
        Platform.__init__(self, pheno, anno, geno, schema, dbtype, dbstring, tmpdir, table_dict, jobs)

class Shard_Illumina(Shard):
    def __init__(self, phenofiles, annotfile, genofile, schema, alleletype, dbtype, dbstring, tmpdir, jobs):
        from dbschema import make_tables_shard_illumina
        table_dict = make_tables_shard_illumina(schema)
        pheno = Pheno_Base(phenofiles, dbstring, table_dict)
        anno = Anno_Illumina(annotfile, schema, tmpdir, table_dict)
        geno = Geno_Shard_Illumina(genofile, schema, dbstring, tmpdir, table_dict, alleletype, jobs)
        Platform.__init__(self, pheno, anno, geno, schema, dbtype, dbstring, tmpdir, table_dict, jobs)

    def load_dataset(self):
        import sys
        alleletype = self.geno.alleletype
        if alleletype != "forward" and alleletype != "top":
            sys.exit("alleletype must be either 'forward' or 'top', not %s"%alleletype)
        super(Shard_Illumina, self).load_dataset()
